{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "\n",
       "require(['notebook/js/codecell'], function(codecell) {\n",
       "  codecell.CodeCell.options_default.highlight_modes[\n",
       "      'magic_text/x-csrc'] = {'reg':[/^%%microblaze/]};\n",
       "  Jupyter.notebook.events.one('kernel_ready.Kernel', function(){\n",
       "      Jupyter.notebook.get_cells().map(function(cell){\n",
       "          if (cell.cell_type == 'code'){ cell.auto_highlight(); } }) ;\n",
       "  });\n",
       "});\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "\n",
    "sys.path.append(os.path.abspath(\"../common\"))\n",
    "\n",
    "import math\n",
    "import time\n",
    "import numpy as np\n",
    "import cv2\n",
    "import pynq\n",
    "import dac_sdc\n",
    "import ctypes\n",
    "\n",
    "team_name = 'BJUT_Runner'\n",
    "team = dac_sdc.Team(team_name, batch_size = 24)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "got nn accelerator!\n"
     ]
    }
   ],
   "source": [
    "last_bais = np.load('last_bias.npy')\n",
    "last_bais = last_bais.reshape((6, 6))[:, :4]\n",
    "cfuns = ctypes.cdll.LoadLibrary(\"./load_image.so\")\n",
    "\n",
    "overlay = pynq.Overlay(team.get_bitstream_path())\n",
    "dma = overlay.axi_dma_0\n",
    "xlnk = pynq.Xlnk()\n",
    "nn_ctrl = overlay.ultra_net_0\n",
    "print('got nn accelerator!')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total time: 3.558751106262207 seconds\n",
      "Total energy: 18.6036615236 J\n",
      "images nums: 1000 .\n",
      "fps: 280.9974539214995 .\n"
     ]
    }
   ],
   "source": [
    "BATCH_SIZE = team.batch_size\n",
    "IMAGE_RAW_ROW = 360\n",
    "IMAGE_RAW_COL = 640\n",
    "IMAGE_ROW = 160\n",
    "IMAGE_COL = 320\n",
    "GRID_ROw = 10\n",
    "GRID_COL = 20\n",
    "X_SCALE = IMAGE_RAW_COL / IMAGE_COL\n",
    "Y_SCALE = IMAGE_RAW_ROW / IMAGE_ROW\n",
    "\n",
    "\n",
    "in_buffer0 = xlnk.cma_array(shape=(BATCH_SIZE, IMAGE_RAW_ROW, IMAGE_RAW_COL, 3), dtype=np.uint8, cacheable = 1)\n",
    "in_buffer1 = xlnk.cma_array(shape=(BATCH_SIZE, IMAGE_RAW_ROW, IMAGE_RAW_COL, 3), dtype=np.uint8, cacheable = 1)\n",
    "in_buffers = [in_buffer0, in_buffer1]\n",
    "out_buffer0 = xlnk.cma_array(shape=(BATCH_SIZE, GRID_ROw, GRID_COL, 6, 6), dtype=np.int32, cacheable = 1)\n",
    "out_buffer1 = xlnk.cma_array(shape=(BATCH_SIZE, GRID_ROw, GRID_COL, 6, 6), dtype=np.int32, cacheable = 1)\n",
    "out_buffers = [out_buffer0, out_buffer1]\n",
    "\n",
    "\n",
    "# use c code load image\n",
    "def load_image(image_paths, buff):\n",
    "    paths = [str(path) for path in image_paths]\n",
    "    tmp = np.asarray(buff)\n",
    "    dataptr = tmp.ctypes.data_as(ctypes.c_char_p)\n",
    "    paths_p_list = [ctypes.c_char_p(bytes(str_, 'utf-8')) for str_ in paths]\n",
    "    paths_c = (ctypes.c_char_p*len(paths_p_list))(*paths_p_list)\n",
    "    cfuns.load_image(paths_c, dataptr, len(paths), IMAGE_ROW, IMAGE_COL, 3)\n",
    "    \n",
    "def sigmoid(x):\n",
    "    s = 1 / (1 + np.exp(-x))\n",
    "    return s\n",
    "\n",
    "def yolo(out_buffer, batch_n, div, last_bais=None, result=None):\n",
    "    res_np = np.array(out_buffer[:batch_n]).reshape(batch_n, -1, 6, 6)\n",
    "    conf = res_np[...,4].sum(axis=2)\n",
    "    max_index = conf.argmax(1)\n",
    "    \n",
    "    grid_x = max_index % GRID_COL\n",
    "    grid_y = max_index // GRID_COL\n",
    "    \n",
    "    boxs = np.zeros((batch_n, 6, 4))\n",
    "    for i in range(batch_n):\n",
    "        boxs[i, :, :] = res_np[i, max_index[i], :, :4] / div + last_bais\n",
    "    xy = sigmoid(boxs[..., :2]).mean(axis=1)\n",
    "    wh = np.exp(boxs[..., 2:4]).mean(axis=1)\n",
    "    \n",
    "    xy[:, 0] += grid_x\n",
    "    xy[:, 1] += grid_y\n",
    "\n",
    "    xy *= 16\n",
    "    wh *= 20\n",
    "\n",
    "    xy[:, 0] *= X_SCALE\n",
    "    xy[:, 1] *= Y_SCALE\n",
    "    wh[:, 0] *= X_SCALE\n",
    "    wh[:, 1] *= Y_SCALE\n",
    "    xmin = xy[:, 0] - wh[:, 0] / 2\n",
    "    xmax = xy[:, 0] + wh[:, 0] / 2\n",
    "    ymin = xy[:, 1] - wh[:, 1] / 2\n",
    "    ymax = xy[:, 1] + wh[:, 1] / 2\n",
    "    \n",
    "    for i in range(batch_n):\n",
    "        temp = [int(xmin[i]), int(xmax[i]), int(ymin[i]), int(ymax[i])]\n",
    "        result.append(temp)\n",
    "\n",
    "which_buffer = 0\n",
    "first_batch = True\n",
    "net_cnt = 0\n",
    "last_batch_size = BATCH_SIZE\n",
    "\n",
    "def net(img_paths, result):\n",
    "    global first_batch\n",
    "    global which_buffer    \n",
    "    global net_cnt\n",
    "    global last_batch_size\n",
    "    # buffer first batch\n",
    "    if first_batch == True:\n",
    "        first_batch = False\n",
    "        which_buffer = 0\n",
    "        load_image(img_paths, in_buffers[which_buffer])\n",
    "        return\n",
    "    # count\n",
    "    net_cnt += 1\n",
    "    nn_ctrl.write(0x0, 0) # Reset\n",
    "    nn_ctrl.write(0x10, in_buffers[which_buffer].shape[0])\n",
    "    nn_ctrl.write(0x0, 1) # Deassert reset\n",
    "    dma.recvchannel.transfer(out_buffers[which_buffer])\n",
    "    dma.sendchannel.transfer(in_buffers[which_buffer])\n",
    "    \n",
    "    # switch buffer\n",
    "    if which_buffer == 0:\n",
    "        which_buffer = 1\n",
    "    else:\n",
    "        which_buffer = 0\n",
    "    # buffer next batch\n",
    "    if img_paths is not None:\n",
    "        load_image(img_paths, in_buffers[which_buffer])\n",
    "    \n",
    "    # yolo \n",
    "    if net_cnt > 1:\n",
    "        yolo(out_buffers[which_buffer], BATCH_SIZE, 7 * 15, last_bais, result)\n",
    "    \n",
    "    if img_paths is not None and len(img_paths) != BATCH_SIZE:\n",
    "        last_batch_size = len(img_paths)\n",
    "            \n",
    "    dma.sendchannel.wait()\n",
    "    dma.recvchannel.wait()\n",
    "    # last batch \n",
    "    if img_paths is None:\n",
    "        yolo(out_buffers[(which_buffer + 1) % 2], last_batch_size, 7 * 15, last_bais, result) \n",
    "        \n",
    "################################Inference##################################\n",
    "interval_time = 0\n",
    "total_time = 0\n",
    "total_energy = 0\n",
    "result = list()\n",
    "team.reset_batch_count()\n",
    "\n",
    "rails = pynq.get_rails()\n",
    "\n",
    "start = time.time()    \n",
    "recorder = pynq.DataRecorder(rails[\"5V\"].power)\n",
    "with recorder.record(0.05): \n",
    "    while True:\n",
    "        image_paths = team.get_next_batch()\n",
    "        net(image_paths, result)\n",
    "\n",
    "#         print('pro_image_cnt', len(result))\n",
    "        # end\n",
    "        if image_paths is None:\n",
    "            break\n",
    "\n",
    "end = time.time()\n",
    "t = end - start\n",
    "    \n",
    "# Energy measurements    \n",
    "energy = recorder.frame[\"5V_power\"].mean() * t    \n",
    "# energy = 0\n",
    "\n",
    "total_time = t\n",
    "total_energy = energy\n",
    "print(\"Total time:\", total_time, \"seconds\")\n",
    "print(\"Total energy:\", total_energy, \"J\")\n",
    "print('images nums: {} .'.format(len(result)))\n",
    "print('fps: {} .'.format(len(result) / total_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1000\n"
     ]
    }
   ],
   "source": [
    "print(len(result))\n",
    "\n",
    "team.save_results_xml(result, total_time, energy)\n",
    "xlnk.xlnk_reset()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
